#!/usr/bin/env python

import os,sys

from optparse import OptionParser

def doGenerate(options):
    #show prompt and ask to confirm if a valid key already exist
    ret=doTest(options)
    if ret==0:
        import gtk
	md = gtk.MessageDialog(parent=None,
			       flags=gtk.DIALOG_DESTROY_WITH_PARENT,
			       type=gtk.MESSAGE_QUESTION,
			       buttons=gtk.BUTTONS_YES_NO,
			       message_format="You have an existing key at %s. Click YES to Overwrite"%options.privatekeyfile)
	response=md.run()
	md.destroy()
	
        if response!=gtk.RESPONSE_YES:
            #return failure
            return -3
        #Remove the existing key files
        try:
            if os.path.exists(options.privatekeyfile):
                os.remove(options.privatekeyfile)
            if os.path.exists(options.privatekeyfile+".pub"):
                os.remove(options.privatekeyfile+".pub")
        except:
            print "Can't remove file %s or %s"%(options.privatekeyfile, options.privatekeyfile+".pub")
            raise
    
    from newpass import doNewPass
    (status,password)=doNewPass()

    if not status:
        return -2

    #call ssh-keygen to generate the key
#    import subprocess
#    cmdstr=["/usr/bin/ssh-keygen","-t","rsa","-P",password,"-f",options.privatekeyfile]
#    retcode=subprocess.call(cmdstr)
    print "Create Key"
    from M2Crypto import RSA
    from base64 import b64encode
    def passwordcallback(*args):
        return password
    rsa=RSA.gen_key(2048,65537,callback=passwordcallback)
    raw_key=rsa.pub()[1]
#    print raw_key
    b64key=b64encode(raw_key)
#    print b64key
    pubkey="ssh-rsa "+b64encode('\x00\x00\x00\x07ssh-rsa%s%s' % (rsa.pub()[0],rsa.pub()[1])) 

#    mem=BIO.MemoryBuffer()
#    rsa.save_pub_key_bio(mem)
#    pubkey=mem.getvalue()
    prikey=rsa.as_pem(callback=passwordcallback)

    #write public key
    ouf=open(options.privatekeyfile+".pub","w")
    ouf.write(pubkey)
    ouf.close()

    #write private key
    ouf=open(options.privatekeyfile,"w")
    ouf.write(prikey)
    ouf.close()
    print "Key created"

    #add the public key to authorized_files
#    inf=open(options.privatekeyfile+".pub","r")
#    pubkey=inf.read()
#    inf.close()

    try:
        #Write key
        ouf=open(options.installpublickey,"a")
        ouf.write(pubkey)
        ouf.write("\n")
        ouf.close()

        import stat
        #change file permission to 0600
        os.chmod(options.installpublickey,stat.S_IRUSR|stat.S_IWUSR)
    except:
        print "Can't install key"
        raise

    return 0
        
def doTest(options):
    #Note: this will only check if the privatekeyfile exist, and permissions are correctly set, DOES NOT check if the file is actually a valid privatekey
    if not os.path.isfile(options.privatekeyfile):
#        print "File doesn't exist"
        return -10

    import stat
    #it should not be rwx-able by group nor other
    mystat=os.stat(options.privatekeyfile).st_mode
    if mystat&stat.S_IRWXG>0 or mystat&stat.S_IRWXO>0:
        print "File Permission Wrong"
        return -11
    
    return 0

def doAuth(options):
#    from askpass import doAskPass
#    (status,password)=doAskPass()
#    if not status:
#        print "password not get"
#        return -21

    import copy
    myenv=copy.deepcopy(os.environ)
    myenv["SSH_ASKPASS"]="%s/askpass.py"%sys.path[0]
    
    import subprocess
    myfd=open("/dev/null","r")
    p=subprocess.Popen(["/usr/bin/ssh-add",options.privatekeyfile],stdin=myfd,env=myenv)
    myfd.close()
    retcode=p.wait()
    
    return 0

def doDelete(options):
    import gtk
    md = gtk.MessageDialog(parent=None,
                           flags=gtk.DIALOG_DESTROY_WITH_PARENT,
                           type=gtk.MESSAGE_QUESTION,
                           buttons=gtk.BUTTONS_YES_NO,
                           message_format="This will remove all credentials from the ssh-agent, you can re-authenticate by choosing \"Authenticate\"")
    response=md.run()
    md.destroy()
    
    if response!=gtk.RESPONSE_YES:
        #return failure
        return -3

    import subprocess
    cmdstr=["/usr/bin/ssh-add","-D"]
    retcode=subprocess.call(cmdstr)

def doRemove(options):
    import gtk
    md = gtk.MessageDialog(parent=None,
                           flags=gtk.DIALOG_DESTROY_WITH_PARENT,
                           type=gtk.MESSAGE_QUESTION,
                           buttons=gtk.BUTTONS_YES_NO,
                           message_format="Are you sure want to remove the key at  %s."%options.privatekeyfile)
    response=md.run()
    md.destroy()
    
    if response!=gtk.RESPONSE_YES:
        #return failure
        return -3
    
    #will remove the public/private key file
    try:
        if os.path.exists(options.privatekeyfile):
            os.remove(options.privatekeyfile)
        if os.path.exists(options.privatekeyfile+".pub"):
            os.remove(options.privatekeyfile+".pub")
    except:
        print "Can't remove file %s or %s"%(options.privatekeyfile, options.privatekeyfile+".pub")
        raise

    import subprocess
    cmdstr=["/usr/bin/ssh-add","-D"]
    retcode=subprocess.call(cmdstr)

    return 0

if __name__ == "__main__":
    
    class MyParser(OptionParser):
        def format_description(self, formatter):
            return self.description
        
    parser = MyParser(description="This program manages the nx-access-key")
        
    parser.add_option("-g", "--generate",action="store_true",dest="generate",help="(Re)generate the nx-access-key, will overwrite any existing one without prompting",default=False)
    parser.add_option("-r", "--remove",action="store_true",dest="remove",help="Destroy and delete the nx-access-key if existed. Otherwise do nothing.",default=False)
    parser.add_option("-d", "--delete",action="store_true",dest="delete",help="Delete the key entry from the current ssh-agent credential list",default=False)    
    
    parser.add_option("-a", "--auth",action="store_true",dest="auth",help="Add the nx-access-key to the current ssh-agent, will return error if nx-access-key does not exist",default=False)

    parser.add_option("-t", "--test",action="store_true",dest="test",help="Test if the nx-access-key exists and permissions are set correctly",default=False)

    parser.add_option("--privatekey-file",type="string",dest="privatekeyfile",help="File that holds the private key, the public key file will be in the same directory, ending with.pub",metavar="PRIVATEKEYFILE",default=os.environ["HOME"]+"/.ssh/nx-access-key")
#    parser.add_option("--publickey-file",type="string",dest="publickeyfile",help="File that holds the public key",metavar="PUBLICKEYFILE",default=os.environ["HOME"]+"/.ssh/nx-access-key.pub")

    parser.add_option("--askpass-msg",type="string",dest="askpassmsg",help="Message to display when showing the askpass window",metavar="ASKPASSMSG",default="Authenticate Key! Please type the passphrase for your key:")
    parser.add_option("--newpass-msg",type="string",dest="newpassmsg",help="Message to display when showing the newpass window",metavar="NEWPASSMSG",default="Create Key! Please type the passphrase for your key twice:")    
    
    parser.add_option("--install-publickey",type="string",dest="installpublickey",help="path/name of the authorized_keys file to install the key after key generation (needs -g), will not install public key if this option is not provided",metavar="INSTALLPUBLICKEY", default=os.environ["HOME"]+"/.ssh/authorized_keys")

#    parser.add_option("--auto-create-dir",action="store_true",dest="autocreatedir",help="Create necessary directories when creating keyfiles (in -g mode)",default=False)
    
    (options, args) = parser.parse_args()

    #Can only set one of the options
    nSet=0
    if options.generate:
        nSet=nSet+1
    if options.remove:
        nSet=nSet+1
    if options.auth:
        nSet=nSet+1
    if options.test:
        nSet=nSet+1
    if options.delete:
        nSet=nSet+1

    if nSet>1:
        print "Have to set one and only one of -t, -a, -d and -g (use --help to see details"
        sys.exit(-1)

    #Check if the options.privatekeyfile and options.publickeyfile are valid
    
    #Case the options and do proper action
    if options.generate:
        sys.exit(doGenerate(options))
    if options.remove:
        sys.exit(doRemove(options))
    if options.auth:
        sys.exit(doAuth(options))
    if options.test:
        sys.exit(doTest(options))
    if options.delete:
        sys.exit(doDelete(options))
        
    #Auto Mode
    if doTest(options)==0:
        sys.exit(doAuth(options))

    #Key not yet properly set, generate
    ret=doGenerate(options)

    #If generation succeed, try auth with it
    if ret==0:
        import gtk
	md = gtk.MessageDialog(parent=None,
			       flags=0,
			       type=gtk.MESSAGE_INFO,
			       buttons=gtk.BUTTONS_OK,
			       message_format="Your key has been created. Press OK to install the key to ssh-agent. You'll need to type your password again.")
	response=md.run()
	md.destroy()
	
        sys.exit(doAuth(options))

    #Otherwise exit with error
    sys.exit(ret)
    
        
        
        
